from motor.motor_asyncio import AsyncIOMotorClient
from pymongo.errors import ConnectionFailure
import os
from dotenv import load_dotenv

# 从.env文件加载环境变量
load_dotenv()

# MongoDB连接详情(使用环境变量)
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017")
DB_NAME = os.getenv("MONGO_DB_NAME", "dbcomparex")

# 创建Motor客户端
client = AsyncIOMotorClient(MONGO_URI)
database = client[DB_NAME]

# 集合
db_info_setting_collection = database.get_collection("db_info_setting")

# 数据同步任务相关集合
task_data_sync_setting_collection = database.get_collection("task_data_sync_setting")
task_data_sync_setting_detail_collection = database.get_collection("task_data_sync_setting_detail")
task_data_sync_ts_mapping_collection = database.get_collection("task_data_sync_ts_mapping")
task_data_sync_execution_log_collection = database.get_collection("task_data_sync_execution_log")
task_target_list_collection = database.get_collection("task_target_list")
task_source_list_collection = database.get_collection("task_source_list")
task_ts_contrast_result_collection = database.get_collection("task_ts_contrast_result")

async def connect_to_mongo():
    try:
        # 验证连接是否有效
        await client.admin.command('ping')
        print("成功连接到MongoDB")
    except ConnectionFailure:
        print("连接MongoDB失败")
        raise

async def close_mongo_connection():
    client.close()
    print("MongoDB连接已关闭")